#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import time
import random
import redis
import unittest


# 顺序消费 有限资源，如proxy 如cookie
class BasicResource(object):
    def __init__(self, name, config):
        self.name = name
        self.config = config

    def get_next(self):
        pass


class RdsResource(BasicResource):
    def __init__(self, name, config):
        super(RdsResource, self).__init__(name, config)
        self.rds = config['rds']
        self.idx_key = config.get('idx_key', '{0}_idx'.format(name))
        self.hash_key = config.get('hash_key', '{0}_hash'.format(name))
        self.lock_prifix = config.get('lock_prifix', 'lock_{0}_'.format(name))
        self.lock_sec = config.get('lock_sec', 0)
        self.lock_step = config.get('lock_step', 0)
        self.max_limit = config.get('max_limit', 180)

    def get_next(self):
        cur_index = self.rds.incr(self.idx_key)
        num = self.rds.hlen(self.hash_key)
        if num == 0:
            return None
        field = cur_index % num
        value = self.rds.hget(self.hash_key, field)
        if self.lock_sec > 0:
            lock_key = '{0}_{1}'.format(self.lock_prifix, value)
            # 随机时间区间，最小lock_sec
            min_sec = self.lock_sec + field * self.lock_step
            max_sec = self.lock_sec + (field + 1) * self.lock_step
            max_sec = max_sec if max_sec < self.max_limit else self.max_limit
            min_sec = min_sec if min_sec < max_sec else max_sec - 3
            rand_sec = random.randint(min_sec, max_sec)
            # 尝试锁定直至成功
            while True:
                cur_tm = time.strftime('%Y-%m-%d %H:%M:%S')
                res = self.rds.set(lock_key, cur_tm, ex=rand_sec, nx=True)
                if res is None:
                    time.sleep(0.5)
                else:
                    return value
        return value


class TestRdsResource(unittest.TestCase):
    def init(self):
        redis_host = '127.0.0.1'
        redis_port = 6379
        redis_password = ''
        self.name = 'resource_test'
        self.rds = redis.Redis(host=redis_host, port=redis_port, db=0, password=redis_password, decode_responses=True)

    def test_next_none(self):
        self.init()
        config = {'rds': self.rds}
        rs = RdsResource(self.name, config)
        self.rds.delete(rs.hash_key)
        next_resource = rs.get_next()
        print('\n\nnext = [{0}]'.format(next_resource))
        self.assertIsNone(next_resource, 'None is expected')

    def test_next_not_none(self):
        self.init()
        config = {'rds': self.rds}
        rs = RdsResource(self.name, config)
        cur_tm = time.strftime('%Y-%m-%d %H:%M:%S')
        self.rds.hset(rs.hash_key, 0, cur_tm)
        next_resource = rs.get_next()
        print('\n\nnext = [{0}]'.format(next_resource))
        self.assertEqual(next_resource, cur_tm)


if __name__ == '__main__':
    unittest.main()
